from Government import *
from Matrix import Matrix
from Agent import *
import time
import random

class Sugarscape(Matrix):

    def __init__(self, shape=None, num_agents = 300, max_sugar=5, sugar_growth_rate=.5,
                 rich_sugar_scale=5, rich_regrowth_scale=5, agent_vision=1, agent_metabolism=4, reproduction=False, tax_rate=.5):
        ''' The sugarscape is a matrix of Locations.
        There are two kinds of Locations: rich and regular;
        rich ones have x times as much sugar and y times as fast a regrowth rate.'''
        Matrix.__init__(self, 51, 51)
        self.shape = shape

        self.max_sugar = max_sugar
        self.sugar_growth_rate = sugar_growth_rate
        self.rich_sugar_scale = rich_sugar_scale
        self.rich_regrowth_scale = rich_regrowth_scale
        self.agent_vision = agent_vision
        self.agent_metabolism = agent_metabolism
        self.tax_rate = tax_rate
        self.agents = []
        self.alpha = -2
        self.num_agents = num_agents
        self.timestamp = 0
        self.reproduction = reproduction
        self.total_wealth = 0
        self.populate_sugarscape()
        self.gov = Government(tax_rate=self.tax_rate, sugarscape=self)

    def reset(self):
        '''resets the whole sugarscape'''
        self.agents = []
        self.gov.tax_rate = self.tax_rate
        Matrix.__init__(self,51,51)
        self.agents = []
        self.populate_sugarscape()
        self.timestamp = 0
        return True

    def populate_sugarscape(self):
        ''' updates the blank sugarscape to have the attributes we want.
        also, populates the agents.'''
        for x in range(self.length):
            for y in range(self.width):
                self.add_attributes(x, y)
        while self.num_agents > len(self.agents):
            x = random.randint(0, self.length - 1)
            y = random.randint(0, self.width - 1)
            self.add_agent(x, y)
        #print self.num_agents, len(self.agents)
        return True

    def add_agent(self, x, y):
        '''adds an agent to the sugarscape'''
        if not self.get_location(x, y).get_has_agent():
            if x % 2 == 0:
                agent = Agent(x, y, self.agent_vision, 1, self.agent_metabolism, self, 'male')
            else:
                agent = Agent(x, y, self.agent_vision, 1, self.agent_metabolism, self, 'female')
            self.agents.append(agent)
            return True
        else:
            return False

    def remove_agent(self, agent):
        '''removes an agent from the sugarscape'''
        self.agents.remove(agent)
        return True

    def add_attributes(self, x, y):
        '''add attributes to the sugarscape.
        allows us to make different shapes'''
        loc = self.get_location(x, y)
        if self.shape=='dogbone':
            if x>20 and x<40 and y>20 and y<40:
                self.set_location_attributes(loc, True)
                return True
        self.set_location_attributes(loc, False)
        return True

    def set_location_attributes(self, loc, rich):
        '''explicitly assigns Location objects values'''
        if rich:
            loc.set_max_sugar(self.max_sugar*self.rich_sugar_scale)
            loc.set_sugar_amt(self.max_sugar*self.rich_sugar_scale)
            loc.set_sugar_growth_rate(self.sugar_growth_rate*self.rich_regrowth_scale)
            loc.is_rich = True
            return True
        else:
            loc.set_max_sugar(self.max_sugar)
            loc.set_sugar_amt(self.max_sugar)
            loc.set_sugar_growth_rate(self.sugar_growth_rate)
            loc.is_rich = False
            return True

    def get_location(self, x, y):
        return self.get_value(x,y)

    def get_nearby_locations(self, vision, location):
        ''' returns a list of all location objects within
        a certain range of the object'''
        nearby_locations = []
        for i in range(-vision, vision + 1):
            if i==0: continue
            nearby_xloc = self.get_location((location.x + i) % self.length, location.y)
            nearby_yloc = self.get_location(location.x, (location.y + i) % self.width)
            nearby_locations.append(nearby_xloc)
            nearby_locations.append(nearby_yloc)
        return nearby_locations

    def get_nearby_agents(self, location):
        ''' returns a list of Agent objects within 1 distance unit from the location'''
        nearby_locations = self.get_nearby_locations(1, location)
        nearby_agents = []
        for agent in self.agents:
            if agent.get_location() in nearby_locations:
                nearby_agents.append(agent)
        return nearby_agents
    
    def nextstep(self):
        '''moves the sugarscape forward one time step.
        Agents and all locations update'''
        for x in range(self.length):
            for y in range(self.width):
                self.get_location(x, y).regrow_sugar()
        self.total_wealth = 0
        for agent in self.agents:
            self.total_wealth += agent.sugar_reserve
        for agent in self.agents:
            agent.nextstep()
        self.timestamp +=1
        if self.timestamp % 1 == 0:
            self.gov.collect_taxes()
            self.gov.redistribute()

    def make_wealth_histogram(self, scale):
        '''makes a histogram with different
        bins'''
        histogram = {}
        for agent in self.agents:
            wealth = int(agent.get_sugar_reserve()/scale)
            wealth_bin = wealth*scale
            histogram[wealth_bin] = histogram.get(wealth_bin, 0) + 1
        return histogram

    def agent_wealths(self):
        ''' makes a list of all agent wealths '''
        wealths = []
        for agent in self.agents:
            wealths.append(agent.get_sugar_reserve())
        return wealths

    def __str__(self):
        return_str = ""
        for y in range(self.width):
            for x in range(self.length):
                return_str += str(self.get_location(x, y).get_sugar_amt())
                return_str += " "
            return_str += "\n"
        return return_str
    
    def set_agent_vision(self, value):
        self.agent_vision = value
        
    def set_tax_rate(self, value):
        self.tax_rate = value
    
    def set_max_sugar(self, value):
        self.max_sugar = value
    
    def set_sugar_growth_rate(self, value):
        self.sugar_growth_rate = value
    
    def set_num_agents(self, value):
        self.num_agents = value
    
    def set_agent_metabolism(self, value):
        self.agent_metabolism = value
        


if __name__=='__main__':
    a = Sugarscape(10,10)
    for i in range(10):
        a.nextstep()
    a.get_location(9, 9).set_sugar_amt(999999)
    print a
    #while True:
    #    print a
     #   a.nextstep()
      #  time.sleep(.1)




